import asyncio
import aiohttp
import json
import pandas as pd
import time
import argparse
from typing import List, Dict
import random

# 配置参数
parser = argparse.ArgumentParser(description='测试服务器并发处理能力')
parser.add_argument('--num_requests', type=int, default=10, help='并发请求数量')
parser.add_argument('--excel_file', type=str, default='laixinjian.xlsx', help='Excel文件路径')
parser.add_argument('--server_url', type=str, default='http://localhost:8000/api/getLetters', help='服务器URL')
args = parser.parse_args()

async def send_request(session: aiohttp.ClientSession, data: Dict) -> Dict:
    """发送单个请求"""
    try:
        async with session.post(args.server_url, json=data) as response:
            return await response.json()
    except Exception as e:
        return {"success": "false", "message": str(e), "data": {}}

async def main():
    # 读取Excel文件
    try:
        df = pd.read_excel(args.excel_file, sheet_name="Sheet1")
    except Exception as e:
        print(f"读取Excel文件失败: {str(e)}")
        return

    # 准备测试数据
    test_data_list = []
    for _, row in df.iterrows():
        test_data = {
            "bh": str(row["信访举报件编号"]),
            "caseIntruction": str(row["原件描述"])
        }
        test_data_list.append(test_data)

    # 如果测试数据不足，则重复使用现有数据
    while len(test_data_list) < args.num_requests:
        test_data_list.extend(test_data_list[:args.num_requests - len(test_data_list)])

    # 随机选择指定数量的测试数据
    selected_data = random.sample(test_data_list, args.num_requests)

    print(f"开始并发测试，请求数量: {args.num_requests}")
    start_time = time.time()

    # 创建会话并发送并发请求
    async with aiohttp.ClientSession() as session:
        tasks = [send_request(session, data) for data in selected_data]
        results = await asyncio.gather(*tasks)

    end_time = time.time()
    total_time = end_time - start_time

    # 统计结果
    success_count = sum(1 for r in results if r.get("success") == "true")
    fail_count = len(results) - success_count

    print("\n测试结果统计:")
    print(f"总请求数: {len(results)}")
    print(f"成功请求数: {success_count}")
    print(f"失败请求数: {fail_count}")
    print(f"总耗时: {total_time:.2f}秒")
    print(f"平均每个请求耗时: {total_time/len(results):.2f}秒")
    print(f"QPS (每秒查询数): {len(results)/total_time:.2f}")

    # 保存详细结果到文件
    result_file = f"concurrent_test_result_{int(time.time())}.json"
    with open(result_file, 'w', encoding='utf-8') as f:
        json.dump({
            "test_config": {
                "num_requests": args.num_requests,
                "server_url": args.server_url,
                "excel_file": args.excel_file
            },
            "summary": {
                "total_requests": len(results),
                "success_count": success_count,
                "fail_count": fail_count,
                "total_time": total_time,
                "avg_time_per_request": total_time/len(results),
                "qps": len(results)/total_time
            },
            "detailed_results": results
        }, f, ensure_ascii=False, indent=2)
    print(f"\n详细结果已保存到: {result_file}")

if __name__ == "__main__":
    asyncio.run(main()) 